package main.java.utils.hxy.thread.thread;

import com.wanma.framework.util.IDate;

import java.util.Collection;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.LongAccumulator;

/**
 * 记录平均进栈出栈的等待时间的LinkedBlockingQueue
 *
 * @author kreo
 * @date 2020-3-22 11:19:56
 */
public class TimeLinkedBlockingQueue<E> extends LinkedBlockingQueue<E> {
    private LongAccumulator avgWaitTime = new LongAccumulator((l, r) -> (l + r) / 2, 0);
    private ConcurrentHashMap<E, Long> timeMap = new ConcurrentHashMap<>();

    public TimeLinkedBlockingQueue() {
        super();
    }

    public TimeLinkedBlockingQueue(int capacity) {
        super(capacity);
    }

    public TimeLinkedBlockingQueue(Collection<? extends E> c) {
        super(c);
    }

    @Override
    public boolean add(E e) {
        boolean result = super.add(e);
        timeMap.put(e, IDate.getNowMillis());
        return result;
    }

    @Override
    public void put(E e) throws InterruptedException {
        super.put(e);
        timeMap.put(e, IDate.getNowMillis());
    }

    @Override
    public boolean offer(E e) {
        boolean result = super.offer(e);
        timeMap.put(e, IDate.getNowMillis());
        return result;
    }

    @Override
    public boolean offer(E e, long timeout, TimeUnit unit) throws InterruptedException {
        boolean result = super.offer(e, timeout, unit);
        timeMap.put(e, IDate.getNowMillis());
        return result;
    }

    @Override
    public boolean remove(Object o) {
        try {
            boolean result = super.remove(o);
            return result;
        } finally {
            calculate((E) o);
        }
    }

    @Override
    public E remove() {
        E e = null;
        try {
            e = super.remove();
            return e;
        } finally {
            calculate(e);
        }
    }

    @Override
    public E element() {
        E e = null;
        try {
            e = super.element();
            return e;
        } finally {
            calculate(e);
        }
    }

    @Override
    public E poll() {
        E e = null;
        try {
            e = super.poll();
            return e;
        } finally {
            calculate(e);
        }
    }

    @Override
    public E poll(long timeout, TimeUnit unit) throws InterruptedException {
        E e = null;
        try {
            e = super.poll(timeout, unit);
            return e;
        } finally {
            calculate(e);
        }
    }

    @Override
    public E take() throws InterruptedException {
        E e = null;
        try {
            e = super.take();
            return e;
        } finally {
            calculate(e);
        }
    }

    @Override
    public E peek() {
        E e = null;
        try {
            e = super.peek();
            return e;
        } finally {
            calculate(e);
        }
    }

    public long getAvgWaitTime() {
        return avgWaitTime.longValue();
    }

    /**
     * 计算时间
     */
    private void calculate(E o) {
        long e = IDate.getNowMillis();
        if (o != null && timeMap.containsKey(o)) {
            Long s = timeMap.get(o);
            timeMap.remove(o);
            avgWaitTime.accumulate(e - s);
        }
    }

    /**
     * 返回平均时间 , 并Reset开始重新计算
     *
     * @return
     */
    public long resetAvgWaitTime() {
        return avgWaitTime.getThenReset();
    }

    public static void main(String[] args) {

    }

}
